'''
KNN 插值
'''

import sys
sys.path.append(".")
sys.path.append('..')

from fancyimpute import KNN
import numpy as np


from impute_method.impute_base import ImputeBase


class KNNImpute(ImputeBase):
    def __init__(self, **kwargs):
        super(KNNImpute, self).__init__(**kwargs)
        # self.kwargs = kwargs
        self.k = kwargs['k']
        self.method = 'knn'

    def impute(self):
        masked_data = self.masked_data[:, 0, :]  # 转化为二维 [n_samples, feature_dim]
        masked_data[masked_data == 0] = np.nan
        imputed_data = KNN(k=self.k).fit_transform(masked_data)
        return self.devide_window(imputed_data)


def main(seq_len, k):
    kwargs = dict()
    kwargs['data_name'] = 'water'
    kwargs['data_type'] = 'masked'
    kwargs['indicators'] = 'WATER_TEMPERATURE,PH_VALUE,TOTAL_NITROGEN,DISSOLVED_OXYGEN'
    kwargs['masked_indicator'] = ''  # 所有指标都有缺失
    kwargs['seq_len'] = seq_len
    kwargs['k'] = k

    knn_impute = KNNImpute(**kwargs)
    knn_impute.metirc()


if __name__ == '__main__':
    # seq_len_list = [24, 48]
    # k_list = [4, 8, 16, 32]
    # for seq_len in seq_len_list:
    #     for k in k_list:
    #         main(seq_len, k)
    main(12, 32)

    # main(24, 16)
    # main(24, 96)

